package mongodb

import (
	"context"
	"time"

	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/mongo/readpref"
)

// Config MongoDB配置
type Config struct {
	URI       string
	Databases map[string]string
}

// Client MongoDB客户端
type Client struct {
	client    *mongo.Client
	databases map[string]*mongo.Database
}

// NewClient 创建MongoDB客户端
func NewClient(config Config) (*Client, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	clientOptions := options.Client().ApplyURI(config.URI)
	client, err := mongo.Connect(ctx, clientOptions)
	if err != nil {
		return nil, err
	}

	// 检查连接
	if err := client.Ping(ctx, readpref.Primary()); err != nil {
		return nil, err
	}

	// 初始化数据库映射
	databases := make(map[string]*mongo.Database)
	for key, dbName := range config.Databases {
		databases[key] = client.Database(dbName)
	}

	return &Client{
		client:    client,
		databases: databases,
	}, nil
}

// GetDatabase 获取指定数据库
func (c *Client) GetDatabase(key string) *mongo.Database {
	return c.databases[key]
}

// GetCollection 获取指定集合
func (c *Client) GetCollection(dbKey, collection string) *mongo.Collection {
	db, ok := c.databases[dbKey]
	if !ok {
		return nil
	}
	return db.Collection(collection)
}

// Close 关闭连接
func (c *Client) Close(ctx context.Context) error {
	return c.client.Disconnect(ctx)
}
